#!/usr/bin/python
# GalacticModel.py

# In alterpars,
# Major support needed for Bochanski mass function
# Additional support for MF/LF choice needed
# In LumFuncBochanski,
# This is already multiplied by bin width, as supplied... how will this interact with finer partitioning of MF and LF?

# Library modules
from time import asctime
from os import linesep
import math
import numpy
#from matplotlib import use as useBackend
#useBackend('Agg')


# Intake parameters
# Units/types noted are the FINAL units/types, after run through alterpars.

# SIMULATION PARAMETERS
import GM_simulationparameters as spars
simpars = {# Parameters directly set by user
           'isacluster' : spars.isacluster,
           'nodes' : spars.nodes,
           'codedir' : spars.codedir,
           'outdir' : spars.outdir,
           'pydir' : spars.pydir,
           'scale' : spars.scale, # pc float
           'galactocentric' : spars.galactocentric, # Not yet implemented
           'Bochanski2010IMFandLF' : spars.Bochanski2010IMFandLF,
           'BoBand' : spars.BoBand,
           'system_or_single' : spars.system_or_single,
           'LensedMag' : spars.LensedMag, # float # Not yet implemented
           'SourceMag' : spars.SourceMag, # float
           'Magnification' : spars.Magnification, # float # Not yet implemented
           'brightlimit' : spars.brightlimit, # int
           'lres' : spars.lres,
           'bres' : spars.bres,
           'dres' : spars.dres,
           'band' : spars.band,
           'c0_V' : spars.c0_V,
           'rho0_V' : spars.rho0_V,
           'c0_I' : spars.c0_I,
           'rho0_I' : spars.rho0_I,
           'c0' : '',
           'rho0' : '',
           'fitcrho' : spars.fitcrho,
           'fixedrho' : spars.fixedrho,
           'fixedc' : spars.fixedc,
           'latitudeweight' : spars.latitudeweight,
           'latitudeweightexternal' : '',
           'latitudeweightinternal' : '',
           'interpolateMv' : spars.interpolateMv,
           'intervals' : spars.intervals,
           'ApparentLimits' : spars.ApparentLimits,
           'interpolateMass' : spars.interpolateMass,
           'maxmassinterval' : spars.maxmassinterval,
           'pencilBeam' : spars.pencilBeam, # boolean
           'pencilBeamPointing' : spars.pencilBeamPointing, # float array
           'timescales' : spars.timescales,
           'NtRequested' : spars.NtRequested,
           'starcluster' : spars.starcluster,
           'bufferfactor' : spars.bufferfactor,
           # Derived parameters
           'ScaleZ' : '', # float pc
           'dlocs' : '', # float index-like array
           'llocs' : '', # float index-like array
           'blocs' : '', # float index-like array
           'delds' : '',
           'dells' : '',
           'delbs' : '',
           'Nptsdlb' : '', # int
           'PtScaled' : '', # float pc
           'PtScalela' : '', # float radians
           'PtScaleba' : '', # float radians
           'timestamp' : '', # str; encodes the time of running
           'codeindex' : -1,
           'NtGiven' : '', # int, total # of timescales given
           'NtMult' : '', # int, # of timescales per integration element
#           '' : spars.,
#           '' : '',
           }

# PHYSICAL PARAMETERS
import GM_physicalparameters as ppars
phypars = {# From Han
           'R0' : ppars.R0, # pc float
           'hR' : ppars.hR, # pc float
           'hdz' : ppars.hdz, # pc float
           'ScaleList' : ppars.ScaleList, # pc float array
           'bulk_vels' : ppars.bulk_vels, #pc/yr float array
           'sigma_vels' : ppars.sigma_vels, #pc/yr float array
           # Derived parameters, Han
           'vanillaScaleList' : '',
           'modifiedScaleList' : '',
           # Binney & Merrifield 1998
           # The luminosity function (if not overwritten by Bochanski)
           'MvList' : ppars.MvList, # mags array
           'LumFuncList' : ppars.LumFuncList, # stars pc^-3 mags^-1 float array
           # Bochanski et al. 2010
           'BoCorrectedValues' : ppars.BoCorrectedValues,
           'BoUncorrectedValues' : ppars.BoUncorrectedValues,
           'BoSystemLF_rband' : ppars.BoSystemLF_rband,
           'BoSingleLF_rband' : ppars.BoSingleLF_rband,
           'BoSystemLF_Jband' : ppars.BoSystemLF_Jband,
           'BoSingleLF_Jband' : ppars.BoSingleLF_Jband,
           'BoSystemMF' : ppars.BoSystemMF,
           'BoSingleMF' : ppars.BoSingleMF,
#           '' : ppars.,
           # Derived parameters, B & M _OR_ Bochanski
           'delMvList' : '',
           'vanillaMvList' : '',
           'vanillaLumFuncList' : '',
           'vanilladelMvList' : '',
           'modifiedMvList' : '',
           'modifiedLumFuncList' : '',
           'modifieddelMvList' : '',

           # Cox 1999
           # The measured angular density of stars
           'bvals' : ppars.bvals, # degrees array
           'MvVals' : ppars.MvVals, # mags array
           'N_bM_V' : ppars.N_bM_V, # float array log10(stars mag**-1 deg**-2)
           'allsky_V' : ppars.allsky_V, # float array log10(stars mag**-1 deg**-2)
           'N_bM_I' : ppars.N_bM_I, # float array log10(stars mag**-1 deg**-2)
           'allsky_I' : ppars.allsky_I, # float array log10(stars mag**-1 deg**-2)
           'N_bM' : '', # float array log10(stars mag**-1 degree**-2)
           'allsky' : '', # float array log10(stars mag**-1 degree**-2)
            # Derived parameters, Cox
           'limitset' : '', # float mags array
           'delMvVals' : '', # float mags array
           'vanillaMvVals' : '', # float mags array
           'vanillalimitset' : '', # float mags array
           'vanilladelMvVals' : '', # float mags array
           'fittingMvVals' : '', # mags array
           'fittinglimitset' : '', # float mags array
           'fittingdelMvVals' : '', # float mags array
           'modifiedMvVals' : '', # float mags array
           'modifiedlimitset' : '', # float mags array
           'modifieddelMvVals' : '', # float mags array
           # Kroupa et al. 1993
           # The mass function
           'masslimits' : ppars.masslimits, # 2d Msun float array
           'massindices' : ppars.massindices, # 1d unitless float array
           'massmults' : ppars.massmults, # 1d float array
           'MvMass' : ppars.MvMass,
           # Derived parameters, Kroupa
           'delMvMass' : '',
           'limitsMvMass' : '',
           'vanillaMvMass' : '', 
           'vanilladelMvMass' : '',
           'vanillalimitsMvMass' : '',
           'modifiedMvMass' : '',
           'modifieddelMvMass' : '',
           'modifiedlimitsMvMass' : '',
           # Derived parameters, Han + Kroupa
           'MvMassScales' : '',
           'modifiedMvMassScales' : '',
           'vanillaMvMassScales' : '',
           # Rotation velocity from Sofue et al.
           'SofueRotV' : ppars.SofueRotV, # 2d [[pc, pc/year],...] array
           # Derived parameters, Sofue
           'modifiedSofueRotV' : '',
           # Location matrices
           'd' : '', # 1d float array pc
           'deltad' : '', # 1d float array pc
           'la' : '', # 1d float array rad
           'deltala' : '', # 1d float array pc
           'ba' : '', # 1d float array rad
           'deltaba' : '', # 1d float array pc
           'z' : '', # 2d float array pc
           'x' : '', # 3d float array pc
           'y' : '', # 3d float array pc
           'R' : '', # 3d float array pc
           'r' : '', # 3d float array pc
           'phi' : '', # 3d float array rad
#           '' : ppars.,
#           '' : '',
           }



# So I was checking off certain changes through this list by putting a comment after each line... I forget what they were, honestly.  The * indicates that they now all draw descriptions of coordinate systems/locations from GM_Locations.  The + may have been ensuring that each plotting function called a backend so that this would work even on clusters.  The OK was... something about how I defined my distances or volume elements, or something?
# Homemade modules
import GM_alterpars # OK
import GM_Locations
#import GM_nonlinear
import GM_distinction # OK*
import plot_distinction # OK+
import GM_LumFunc # OK
import GM_LumFuncBochanski # OK
import plot_LumFunc # OK+
import GM_fitc0rho0 # OK
# These modules are all secretly used by GM_fitc0rho0
#import GM_apparent # OK
#import GM_logN # OK
#import GM_N_bvals # OK
# Plot the aftermath, then delete the unloved
import plot_fitc0rho0 # OK+
import plot_apparent # OK+*
import plot_logN # OK+
import plot_logN_HanStyle # OK+

import GM_darklight #*
# Secretly used by GM_darklight and GM_MassDistrib and, well, lots of things
#import GM_InterpRoutines
import plot_darklight
import GM_MassDistrib # OK?
import plot_MassDistrib # OK?+

import GM_velocities # OK?
import plot_velocities # OK?+
import GM_EstimateDependence #*?
import plot_EstimateDependence
import GM_SimpleIntegrate # OK?
import plot_SimpleIntegrate # OK?+
import plot_SimpleIntegrateCluster # OK?+


# Begin
print 'Beginning.'
print asctime()
print linesep

# Update parameters so that their units and datatypes are appropriate
(simpars, phypars) = GM_alterpars.alterpars(simpars, phypars)
del GM_alterpars
#simpars = GM_nonlinear.nonlinear(simpars)
#del GM_nonlinear
phypars = GM_Locations.Locations(simpars, phypars)



# Determine distance modulus and extinction for all locations
# DistMod is 1d (d only); ExtMod is 2d (d, b)
(DistMod, ExtMod) = GM_distinction.distinction(simpars, phypars)
del GM_distinction
# Plot a slice through that
#plot_distinction.plot(DistMod, ExtMod, simpars)
del plot_distinction
print 'I have calculated the extinction coefficient and distance modulus.'
print asctime()
print linesep

# Get distribution of stars
# elements of distriblist are 3d (d, l, b)
if simpars['Bochanski2010IMFandLF'] == False:
    distriblist = GM_LumFunc.lumfunc(simpars, phypars)
if simpars['Bochanski2010IMFandLF'] == True:
    distriblist = GM_LumFuncBochanski.lumfunc(simpars, phypars)
del GM_LumFunc
del GM_LumFuncBochanski
if simpars['pencilBeam'] == False:
    # Plot a slice through that
    plot_LumFunc.plot(simpars, phypars, distriblist)
    pass
del plot_LumFunc
print 'I have calculated the distribution of stars.'
print asctime()
print linesep

# Choose rho0 and c0, if necessary
# Warning! This portion of the program is slow! 1 iteration takes ~15 minutes for a 100x100x100 grid with resolution in Mv of one.
if simpars['fitcrho'] == True:
    (simpars, difflist, distlist, distarraylist, paramlist, clocslist, rlocslist) = GM_fitc0rho0.fit(simpars, phypars, DistMod, ExtMod, distriblist)
    print "I have done my best to fit the values of c0 and rho0."
    print asctime()
    print linesep
else:
    print "I will use the values of c0 and rho0 indicated in the parameters file without attempting to find a better fit."
    print linesep

# Calculate everything for chosen rho0 and c0
(appearance, N, allskyN, compN) = GM_fitc0rho0.iterateverbose(simpars, phypars, DistMod, ExtMod, distriblist)
del GM_fitc0rho0
del DistMod, ExtMod, distriblist
print 'I calculated the apparent distribution of stars.'
print asctime()
print linesep

# Plot diagnostics of the fitting process, if necessary
if simpars['fitcrho'] == True:
    plot_fitc0rho0.plot(simpars, phypars, difflist, distlist, distarraylist, paramlist, clocslist, rlocslist)
    del plot_fitc0rho0
    del difflist, distlist, distarraylist, paramlist
    print 'I have plotted diagnostics of the fitting process.'
else:
    print "As c0 and rho0 were not fitted, there are no fitting diagnostics to plot."
    print linesep

if simpars['pencilBeam'] == False:
    # Plot apparent distribution of stars
    plot_apparent.plot(simpars, phypars, appearance)
    # Plot apparent distribution of stars per degree squared
#    plot_logN.plot(simpars, phypars, N, allskyN)
    #...and throw a Han-style comparison in for good measure
#    plot_logN_HanStyle.plot(simpars, phypars, compN, N, allskyN)
    pass
del plot_apparent
del plot_logN
del plot_logN_HanStyle
print 'I plotted the apparent distribution of stars.'
print asctime()
print linesep

#(Mv_comp_vals, BMlist, Klist) = GM_darklight.comparable(simpars, phypars)
# Warning: I don't think I ever actually used the line of code immediately below so it's all untested etc.
#mass0 = GM_darklight.fit(Mv_comp_vals, BMlist, Klist)
del GM_darklight
#plot_darklight.plot(Mv_comp_vals, BMlist, Klist, 'complex')
del plot_darklight
#print 'I have compared the distributions of compact masses and bright stars.'
#print asctime()
#print linesep
# In fact, the entire above section is made of iffy stuff.

# elements of darkstars are 3d (d, l, b)
masstuple = GM_MassDistrib.MassDistrib(simpars, phypars)
#(darkstars, rho0mlist) = masstuple
del GM_MassDistrib
#plot_MassDistrib.plot(masstuple, simpars, phypars)
del plot_MassDistrib
print 'I have created distributions of compact masses.'
print asctime()
print linesep

vtupleMS = GM_velocities.velocities(simpars, phypars)
#(vrelappl, sigmaappl2, sigmaappb2) = vtupleMS
del GM_velocities
if simpars['pencilBeam'] == False:
#    plot_velocities.plot(simpars, phypars, vtuple)
    pass
del plot_velocities
print 'I have computed velocity distributions.'
print asctime()
print linesep

#EDtuple = GM_EstimateDependence.ED(simpars, phypars, appearance, vtuple)
#(numberlist, distancelist, vtapplist, proptolist) = EDtuple
del GM_EstimateDependence
#plot_EstimateDependence.plot(simpars, phypars, EDtuple)
del plot_EstimateDependence


if simpars['isacluster'] == False:
    MSratestuple = GM_SimpleIntegrate.simple(simpars, phypars, appearance, masstuple, vtupleMS, vtupleMS)
    #(lensrates, MvHist, MassHist, DSHist, DLHist, bHist, chosenTimes) = ratestuple
#    if simpars['starcluster'] == True:
        

    plot_SimpleIntegrate.plot(simpars, phypars, MSratestuple, False)
    plot_SimpleIntegrate.plot(simpars, phypars, MSratestuple, True)

    del GM_SimpleIntegrate
    del plot_SimpleIntegrate



if simpars['isacluster'] == True:
    print "Wait up!  Quit the program!  This part is not ready!"
    import cPickle as cP
    from os import mkdir
    from os import listdir
    from time import sleep
    import wrapbeospawn
    foldername = simpars['outdir']+"TempData"+simpars['timestamp']+"/"
    infolder = foldername+"In/"
    outfolder = foldername+"Out/"
    mkdir(foldername)
    mkdir(infolder)
    mkdir(outfolder)
    safe_appearance = list(appearance)
#    safe_simpars = dict(simpars)
    safe_phypars = dict(phypars)
    ncodes = len(safe_appearance)

### For testing only!
#    ncodes = 2
### For testing only!

    m = len(str(ncodes))
    parfilenames = []
    outfilenames = []
    for i in range(0, ncodes):
        appearance = [numpy.array(safe_appearance[i])]
        phypars['MvVals'] = numpy.array([safe_phypars['MvVals'][i]])
        phypars['delMvVals'] = numpy.array([safe_phypars['delMvVals'][i]])
        phypars['limitset'] = numpy.array([safe_phypars['limitset'][i]])
        simpars['codeindex'] = int(i)
#        simpars[''] = safe_simpars[''][i]
        datadict = {'simpars' : simpars,
                    'phypars' : phypars,
                    'appearance' : appearance,
                    'masstuple' : masstuple,
                    'vtuple' : vtuple}
        strindex = str(i)
        n = len(strindex)
        if n < m:
            strindex = "0"*(m-n) + strindex
        parfilename = infolder+"Pars"+strindex+".txt"
        outfilename = outfolder+"Out"+strindex+".txt"
        parfilenames.append(parfilename)
        outfilenames.append(outfilename)
        parfile = open(parfilename, 'w')
        cP.dump(datadict, parfile)
        parfile.close()
        comstr = simpars['pydir']+" "+simpars['codedir']+"GM_SimpleIntegrate.py "+parfilename+" "+outfilename
        wrapbeospawn.wrapbeospawn(comstr, simpars['nodes'])
    del appearance, masstuple, vtuple
    del GM_SimpleIntegrate
    print "Jobs have been distributed to threads."
    print asctime()
    print linesep
    # Every thirty seconds, check for completion; notify
    numdone = 0
    while True:
        currentfiles = listdir(outfolder)
        if numdone < ncodes:
            if len(currentfiles) > numdone:
                numdone = len(currentfiles)
                print str(numdone), "modules have completed and written to file!"
                print asctime()
                print linesep
        if len(currentfiles) == ncodes:
            print "The final module has completed and written to file!"
            print asctime()
            print linesep
            break
        sleep(30)
    # Allow a few seconds to be sure that you are not competing for the use of a file; read back in
    sleep(10)
    appearance = list(safe_appearance)
#    simpars = dict(safe_simpars)
    phypars = dict(safe_phypars)
    simpars['codeindex'] = -1
    #(lensrates, MvHist, MassHist, DSHist, DLHist, bHist) = ratestuple
#    lensrates
    # Read in and concatenate the results. Mv handled differently here.
    MvHistFull = numpy.zeros(len(phypars['MvVals']))
    MassList = []
    DSList = []
    DLList = []
    bList = []    
    for i in range(0, len(outfilenames)):
        outfilei = open(outfilenames[i])
        outpars = cP.load(outfilei)
        outfilei.close()
        theserates = outpars['ratestuple']
        (tlensrates, tMvHist, tMassHist, tDSHist, tDLHist, tbHist) = theserates
        MvHistFull[i] = float(tMvHist[-1]) # Wicked little quirk; quick fix
        MassList.append(tMassHist[-1])
        DSList.append(tDSHist[-1])
        DLList.append(tDLHist[-1])
        bList.append(tbHist[0])
    del theserates, tlensrates, tMvHist, tMassHist, tDSHist, tDLHist, tbHist
    # Arrayify 'em
    MassHistFull = numpy.array(MassList) # 2d
    DSHistFull = numpy.array(DSList) # 2d
    DLHistFull = numpy.array(DLList) # 2d
    bHistFull = numpy.array(bList) # 2d oversize
    # Bring back less-informative versions
    MvValMask = numpy.less_equal(phypars['MvVals'][numpy.newaxis,:],simpars['SourceMag'][:,numpy.newaxis])
    del MassList, DSList, DLList, bList #MvList,
    MvList = []
    MassList = []
    DSList = []
    DLList = []
    bList = []    
    for i in range(0, len(simpars['SourceMag'])):
        MaskedMv = numpy.multiply(MvValMask[i][:], MvHistFull)
        MaskedMass = numpy.multiply(MvValMask[i][:,numpy.newaxis], MassHistFull)
        MaskedDS = numpy.multiply(MvValMask[i][:,numpy.newaxis], DSHistFull)
        MaskedDL = numpy.multiply(MvValMask[i][:,numpy.newaxis], DLHistFull)
        Maskedb = numpy.multiply(MvValMask[i][:,numpy.newaxis], bHistFull)
#        Masked = numpy.multiply(MvValMask[i][:,numpy.newaxis], HistFull)
        ReducedMv = numpy.array(MaskedMv)
        ReducedMass = numpy.add.reduce(MaskedMass, axis=0)
        ReducedDS = numpy.add.reduce(MaskedDS, axis=0)
        ReducedDL = numpy.add.reduce(MaskedDL, axis=0)
        Reducedb = numpy.add.reduce(Maskedb, axis=0)
#        Reduced = numpy.add.reduce(Masked, axis=0)
        MvList.append(ReducedMv)
        MassList.append(ReducedMass)
        DSList.append(ReducedDS)
        DLList.append(ReducedDL)
        bList.append(Reducedb)
#        List.append(Reduced)
    MvHist = numpy.array(MvList)
    MassHist = numpy.array(MassList)
    DSHist = numpy.array(DSList)
    DLHist = numpy.array(DLList)
    bHist = numpy.array(bList)
#    Hist = numpy.array(List)
    del MaskedMv, MaskedMass, MaskedDS, MaskedDL, Maskedb
    del ReducedMv, ReducedMass, ReducedDS, ReducedDL, Reducedb
    del MvList, MassList, DSList, DLList, bList
    # Flatten out MvHist to get lensrates
    lensrates = numpy.add.reduce(MvHist, axis=1)
    lensrates = lensrates.flatten()
    ratestuple = (lensrates, MvHist, MassHist, DSHist, DLHist, bHist)
    # Apply a similar process to get the full cumulative map of bHist
    # Outer index is magnitude, from lowest (brightest, least exclusive) to highest (dimmest, most exclusive)
    # so accumulate naturally along this axis
    bHistCumulative = numpy.add.accumulate(bHistFull, axis=0)
    # Inner index is latitude, from lowest (richest, but most exclusive) to highest (poorest, but least exclusive)
    bHistCumulative = numpy.add.accumulate(bHistFull[:,::-1], axis=1)
    CumulativeAngle = simpars['blocs'][:]*simpars['PtScaleba'] # radians
    CumulativeAreaOld = 4*math.pi*(1 - numpy.cos(CumulativeAngle)) # steradians
    CumulativeAngle = CumulativeAngle*(180./math.pi) # degrees
    CumulativeAreaOld = CumulativeAreaOld*(180./math.pi)**2 # sq. degrees
    CumulativeAreaNew = numpy.arange(len(CumulativeAreaOld))
    CumulativeAreaNew = CumulativeAreaNew/float(CumulativeAreaNew[-1])
    CumulativeAreaNew = CumulativeAreaNew*CumulativeAreaOld[-1]
    bHistCumulativeAreaOld = numpy.array(bHistCumulative.T)
    import GM_InterpRoutines as IR
    bHistCumulativeAreaNew = IR.OldToNewLinear(CumulativeAreaOld, bHistCumulativeAreaOld, CumulativeAreaNew)
    del IR
    bHistCumulativeAreaNew = numpy.array(bHistCumulativeAreaNew.T)
    bHistTuple = (bHistCumulative, CumulativeAngle, bHistCumulativeAreaNew, CumulativeAreaNew)

    # Pickle all the best results, so that in case something goes wrong with plotting, (or already something went wrong, or I get curious later,) they can be recovered
    resultsDictionary = {'MvHistFull' : MvHistFull,
                         'MassHistFull' : MassHistFull,
                         'DSHistFull' : DSHistFull,
                         'DLHistFull' : DLHistFull,
                         'bHistFull' : bHistFull,
                         'MvHist' : MvHist,
                         'MassHist' : MassHist,
                         'DSHist' : DSHist,
                         'DLHist' : DLHist,
                         'bHist' : bHist,
                         'bHistCumulative' : bHistCumulative,
                         'CumulativeAngle' : CumulativeAngle,
                         'bHistCumulativeAreaOld' : bHistCumulativeAreaOld,
                         'CumulativeAreaOld' : CumulativeAreaOld,
                         'bHistCumulativeAreaNew' : bHistCumulativeAreaNew,
                         'CumulativeAreaNew' : CumulativeAreaNew,
#                         '' : ,
                         }

    outfilename = simpars['outdir']+"FinalTally"+simpars['timestamp']
    outfile = open(outfilename, 'w')
    cP.dump(resultsDictionary, outfile)
    outfile.close()
    del cP

    plot_SimpleIntegrate.plot(simpars, phypars, ratestuple, False)
    plot_SimpleIntegrate.plot(simpars, phypars, ratestuple, True)
    del plot_SimpleIntegrate
    plot_SimpleIntegrateCluster.plot(simpars, phypars, bHistTuple)
    del plot_SimpleIntegrateCluster

    del bHistCumulative, CumulativeAngle, bHistCumulativeAreaOld, CumulativeAreaOld, bHistCumulativeAreaNew, CumulativeAreaNew
    del MvHist, MassHist, DSHist, DLHist, bHist
    del MvHistFull, MassHistFull, DSHistFull, DLHistFull, bHistFull




print '''That's all for now!'''
print asctime()
print linesep

